import ksql.utils as utils

from ksql import KSQLAPI
from confluent_kafka import Producer

from cc.experimental.wliang.ksqldb.ksqldb_strategy_log_schema import KSQL_STREAM_CREATE_TEMP
from cc.experimental.wliang.ksqldb.ksqldb_strategy_log_schema import KSQL_STREAM_SCHEMA_DESC


class KSqldbClient(object):
  def __init__(self, ksqldb_client, kafka_topic):
    self.api_cli = KSQLAPI(ksqldb_client, check_version=False)
    self.kafka_topic = kafka_topic
    self.stream_name = kafka_topic.replace(".", "_").replace("-", "_")
    self.schema_desc = KSQL_STREAM_SCHEMA_DESC
    self.format_type = "json"

  def set_kafka_topic(self, kafka_topic):
    self.kafka_topic = kafka_topic

  def set_stream_name(self, stream_name):
    self.stream_name = stream_name

  def get_kafka_topic(self):
    return self.kafka_topic

  def get_stream_name(self):
    return self.stream_name

  def execute(self, sql):
    r = self.api_cli.ksql(sql)
    return r[0]["commandStatus"]["status"] == "SUCCESS"

  def execute_no_throw(self, sql):
    ret = True
    try:
      r = self.api_cli.ksql(sql)
      ret = (r[0]["commandStatus"]["status"] == "SUCCESS")
    except:
      # import traceback
      # traceback.print_exc()
      ret = False
    return ret

  # def initial(self):
  #   # create stream
  #   if self.stream_name.upper() in utils.get_all_streams(self.api_cli):
  #     return
  #   ksql_create_stream = KSQL_STREAM_CREATE_TEMP.format(self.stream_name,
  #                                                       self.schema_desc,
  #                                                       self.kafka_topic,
  #                                                       self.format_type)
  #   if not self.execute(ksql_create_stream):
  #     raise "failed to create stream"

  def initial(self):
    ksql_create_stream = KSQL_STREAM_CREATE_TEMP.format(self.stream_name,
                                                        self.schema_desc,
                                                        self.kafka_topic,
                                                        self.format_type)
    if self.execute_no_throw(ksql_create_stream):
      return self.stream_name
    else:
      return None

  def consume(self, ksql, process_func = None):
    rows = self.api_cli.query(ksql)
    if process_func is None:
      for item in rows: print(item)
    else:
      process_func(rows, self.kafka_topic)


class KafkaClient(object):
  def __init__(self, bootstrap_server, kafka_topic=None):
    self.producer = Producer({"bootstrap.servers": bootstrap_server})
    self.kafka_topic = kafka_topic

  def set_kafka_topic(self, kafka_topic):
    self.kafka_topic = kafka_topic

  def get_kafka_topic(self):
    return self.kafka_topic

  def produce(self, rows, kafka_topic_specify = None):
    kafka_topic = self.kafka_topic if kafka_topic_specify is None else kafka_topic_specify
    if kafka_topic is None:
      raise "kafka topis is not specify"
    for row in rows:
      self.producer.produce(kafka_topic, row)
      self.producer.flush()

